package com.efeichong.javac;

import com.efeichong.lock.EnhanceReentrantLock;
import com.efeichong.lock.Lock;
import com.efeichong.lock.LockType;
import com.efeichong.lock.RedisReentrantLock;
import com.google.auto.service.AutoService;
import com.sun.source.util.TreePath;
import com.sun.tools.javac.api.JavacTrees;
import com.sun.tools.javac.code.Flags;
import com.sun.tools.javac.model.JavacElements;
import com.sun.tools.javac.processing.JavacProcessingEnvironment;
import com.sun.tools.javac.tree.JCTree;
import com.sun.tools.javac.tree.TreeMaker;
import com.sun.tools.javac.util.Context;
import com.sun.tools.javac.util.List;
import com.sun.tools.javac.util.Names;

import javax.annotation.processing.*;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.Element;
import javax.lang.model.element.TypeElement;
import javax.tools.Diagnostic;
import java.util.ArrayList;
import java.util.LinkedHashSet;
import java.util.Set;
import java.util.regex.Matcher;
import java.util.regex.Pattern;


/**
 * @author lxk
 * @date 2020/02/04
 * @description 基于 {@link Lock}注解编译代码实现逻方法锁功能
 */
@SupportedSourceVersion(SourceVersion.RELEASE_8)
@AutoService(Processor.class)
public class LockProcessor extends AbstractProcessor {
    private static final Pattern pattern = Pattern.compile("\\$\\{[a-zA-Z0-9]+(\\.[a-zA-Z0-9]+)+\\(\\)\\}");
    /**
     * 打印编译器log
     */
    private Messager messager;
    /**
     * 待处理的抽象语法树
     */
    private JavacTrees javacTrees;
    /**
     * 创建AST节点的功能
     */
    private TreeMaker treeMaker;
    /**
     * 创建标识符的功能
     */
    private Names names;
    /**
     * 用来处理Element的工具类
     * Elements接口的对象，用于操作元素的工具类。
     */
    private JavacElements elementUtils;

    /**
     * 支持编译时注解类型
     *
     * @return
     */
    @Override
    public Set<String> getSupportedAnnotationTypes() {
        Set<String> annotations = new LinkedHashSet<>();
        annotations.add(Lock.class.getCanonicalName());
        return annotations;
    }

    @Override
    public synchronized void init(ProcessingEnvironment processingEnv) {
        super.init(processingEnv);
        this.messager = processingEnv.getMessager();
        this.javacTrees = JavacTrees.instance(processingEnv);
        Context context = ((JavacProcessingEnvironment) processingEnv).getContext();
        this.treeMaker = TreeMaker.instance(context);
        this.names = Names.instance(context);
        elementUtils = (JavacElements) processingEnv.getElementUtils();
    }

    @Override
    public boolean process(Set<? extends TypeElement> annotations, RoundEnvironment roundEnv) {
        try {
            int lockSuffixIndex = 1000;
            for (Element element : roundEnv.getElementsAnnotatedWith(Lock.class)) {
                Lock lock = element.getAnnotation(Lock.class);
                //锁类型
                String lockTypeName;
                //上锁方法参数类型列表
                java.util.List<JCTree.JCExpression> lockMethodArgTypes = new ArrayList();
                //上锁方法参数值列表
                java.util.List<JCTree.JCExpression> lockMethodArgValues = new ArrayList();
                //解锁方法参数类型列表
                java.util.List<JCTree.JCExpression> unlockMethodArgTypes = new ArrayList();
                //解锁方法参数值列表
                java.util.List<JCTree.JCExpression> unlockMethodArgValues = new ArrayList();

                if (lock.lockType() == LockType.THREAD_LOCK) {
                    Class<EnhanceReentrantLock> reentrantLockClass = EnhanceReentrantLock.class;
                    lockTypeName = reentrantLockClass.getSimpleName();
                    String packageName = reentrantLockClass.getPackage().getName();
                    addImport(element, packageName, lockTypeName);
                } else {
                    Class<RedisReentrantLock> redisReentrantLockClass = RedisReentrantLock.class;
                    lockTypeName = redisReentrantLockClass.getSimpleName();
                    String packageName = redisReentrantLockClass.getPackage().getName();
                    addImport(element, packageName, lockTypeName);

                }
                java.util.List<JCTree.JCStatement> jcStatements = new ArrayList<>();
                String key = lock.key();
                Matcher matcher = pattern.matcher(key);
                JCTree.JCVariableDecl declareKey = null;
                if (matcher.matches()) {
                    String methodPath = key.substring(2, key.length() - 1);
                    String classPath = methodPath.substring(0, methodPath.lastIndexOf("."));
                    String packageName = classPath.substring(0, classPath.lastIndexOf("."));
                    String className = classPath.substring(classPath.lastIndexOf(".") + 1);
                    String methodName = methodPath.substring(methodPath.lastIndexOf(".") + 1).replace("()", "");
                    addImport(element, packageName, className);

                    declareKey = addMethodObject("String", "key", className, methodName);

                    // lock/unlock 方法入参  key
                    lockMethodArgTypes.add(treeMaker.Ident(names.fromString("String")));
                    lockMethodArgValues.add(treeMaker.Ident(names.fromString("key")));

                    unlockMethodArgTypes.add(treeMaker.Ident(names.fromString("String")));
                    unlockMethodArgValues.add(treeMaker.Ident(names.fromString("key")));
                } else {
                    // lock/unlock 方法入参  key
                    lockMethodArgTypes.add(treeMaker.Ident(names.fromString("String")));
                    lockMethodArgValues.add(treeMaker.Literal(key));

                    unlockMethodArgTypes.add(treeMaker.Ident(names.fromString("String")));
                    unlockMethodArgValues.add(treeMaker.Literal(key));
                }


                //lock 方法入参 过期时间
                lockMethodArgTypes.add(treeMaker.Ident(names.fromString("Long")));
                lockMethodArgValues.add(treeMaker.Literal(lock.lockExpireTime()));
                //lock 方法入参 重试次数
                lockMethodArgTypes.add(treeMaker.Ident(names.fromString("Integer")));
                lockMethodArgValues.add(treeMaker.Literal(lock.retryCount()));
                //lock 方法入参 重试时间间隔
                lockMethodArgTypes.add(treeMaker.Ident(names.fromString("Long")));
                lockMethodArgValues.add(treeMaker.Literal(lock.retryTimeInterval()));

                JCTree jcTree = elementUtils.getTree(element);
                JCTree.JCMethodDecl jcMethodDecl = (JCTree.JCMethodDecl) jcTree;
                //添加类中锁的成员变量
                addClassLockObject(element, lockTypeName, lockSuffixIndex);
                //执行lock.lock()方法
                JCTree.JCExpressionStatement lockState = treeMaker.Exec(treeMaker.Apply(List.from(lockMethodArgTypes),
                        treeMaker.Select(
                                //锁实例 lock
                                treeMaker.Ident(names.fromString("$lock" + lockSuffixIndex)),
                                //调用lock()方法
                                names.fromString("lock")
                        ),
                        List.from(lockMethodArgValues)));

                //执行lock.unlock()方法
                JCTree.JCExpressionStatement unLock = treeMaker.Exec(treeMaker.Apply(List.from(unlockMethodArgTypes),
                        treeMaker.Select(
                                //锁实例 lock
                                treeMaker.Ident(names.fromString("$lock" + lockSuffixIndex)),
                                //调用unlock()方法
                                names.fromString("unlock")
                        ),
                        List.from(unlockMethodArgValues)));
                jcMethodDecl.getBody().getStatements().forEach(jcStatement -> {
                    jcStatements.add(jcStatement);
                });
                jcMethodDecl.body.stats = List.from(jcStatements);
                //将原本的方法体用try-finally包裹起来
                JCTree.JCBlock finallyBlock = treeMaker.Block(0, List.of(unLock));
                jcMethodDecl.body = treeMaker.Block(0, List.of(
                        treeMaker.Try(jcMethodDecl.body, List.nil(), finallyBlock)
                        )
                );
                jcMethodDecl.body.stats = jcMethodDecl.body.stats.prepend(lockState);
                if (declareKey != null) {
                    jcMethodDecl.body.stats = jcMethodDecl.body.stats.prepend(declareKey);
                }
                lockSuffixIndex++;
            }
        } catch (Exception e) {
            messager.printMessage(Diagnostic.Kind.ERROR, e.getMessage());
            return false;
        }
        return true;
    }

    /**
     * 添加类中锁的成员变量
     *
     * @param element      程序元素
     * @param lockTypeName 锁类型
     * @param lockSuffix   锁对象后缀名
     */
    private void addClassLockObject(Element element, String lockTypeName, int lockSuffix) {
        JCTree.JCClassDecl classDecl = (JCTree.JCClassDecl) javacTrees.getPath(element).getParentPath().getLeaf();
        List<JCTree> classDefs = classDecl.defs;
        //创建重入锁对象
        JCTree.JCNewClass lockObject = treeMaker.NewClass(
                null,
                List.nil(),
                treeMaker.Ident(names.fromString(lockTypeName)),
                List.nil(),
                null
        );
        //将重入锁对象指向lock实例
        JCTree.JCVariableDecl lockInstance = treeMaker.VarDef(
                treeMaker.Modifiers(Flags.PARAMETER),
                //锁实例 lock
                names.fromString("$lock" + lockSuffix),
                treeMaker.Ident(names.fromString(lockTypeName)),
                lockObject
        );
        classDecl.defs = classDefs.prepend(lockInstance);
    }

    /**
     * 向方法中添加局部变量  构建锁key
     *
     * @param declareClassName 声明对象类型
     * @param instanceName     实例名
     * @param className        静态类
     * @param methodName       静态方法
     * @return
     */
    private JCTree.JCVariableDecl addMethodObject(String declareClassName, String instanceName, String className, String methodName) {
        //调用方法
        JCTree.JCMethodInvocation call = treeMaker.Apply(
                List.nil(),
                treeMaker.Select(
                        treeMaker.Ident(names.fromString(className)),
                        names.fromString(methodName)
                ),
                List.nil()
        );
        //将调用方法的返回指向实例
        JCTree.JCVariableDecl keyDecl = treeMaker.VarDef(
                treeMaker.Modifiers(Flags.PARAMETER),
                names.fromString(instanceName),
                treeMaker.Ident(names.fromString(declareClassName)),
                call
        );
        return keyDecl;
    }

    /**
     * 导入需要的包
     *
     * @param element     程序元素
     * @param packageName 包路径
     * @param className   类名
     */
    private void addImport(Element element, String packageName, String className) {
        TreePath treePath = javacTrees.getPath(element);
        JCTree.JCCompilationUnit jccu = (JCTree.JCCompilationUnit) treePath.getCompilationUnit();
        java.util.List<JCTree> trees = new ArrayList<>();
        trees.addAll(jccu.defs);
        JCTree.JCImport whereImport = buildImport(packageName, className);
        if (!trees.contains(whereImport)) {
            trees.add(whereImport);
        }
        jccu.defs = List.from(trees);
    }

    /**
     * 构建导入
     *
     * @param packageName 包路径
     * @param className   类名
     * @return
     */
    private JCTree.JCImport buildImport(String packageName, String className) {
        JCTree.JCIdent ident = treeMaker.Ident(names.fromString(packageName));
        JCTree.JCImport jcImport = treeMaker.Import(treeMaker.Select(
                ident, names.fromString(className)), false);
        return jcImport;
    }

}